#!/usr/bin/env python3

import sys
import random
import math
import hashlib
import struct

nr_node = 6
nr_ssd_per_node = 6
nr_pg = 1024

nr_replica = 3

g_cluster_id = ''

# 
# {node_name: [osd_name, xxxx, xxxx], ...}
#
g_node_list = {}

g_node_weight = {}
g_osd_weight = {}

def get_osd_name(osd_id):
	osd_name = 'o_%02x_%016x' % (osd_id, random.randint(0, 0xffffffffffffffff))
	return osd_name

def get_node_name(node_id):
	node_name = 'n_%02x_%016x' % (node_id, random.randint(0, 0xffffffffffffffff))
	return node_name

def sys_init():
	global g_node_list, g_cluster_id
	g_node_list = {}
	g_cluster_id = '%016x' % random.randint(0, 0xffffffffffffffff)	
	#print(g_cluster_id)
	osd_id = 0
	for i in range(nr_node):
		osd_list = []
		for j in range(nr_ssd_per_node):
			osd_name = get_osd_name(osd_id)
			osd_list.append(osd_name)
			osd_id += 1
			#print(i, j, osd_name)
		node_name = get_node_name(i)
		g_node_list[node_name] = osd_list
	#print(g_node_list)
	return

def add_disk_by_name(node_name):
	osd_idx = len(g_node_list[node_name])
	osd_name = get_osd_name(osd_idx)
	g_node_list[node_name].append(osd_name)
	return osd_name

def add_disk(node_id):
	global g_node_list
	node_name_list = list(g_node_list.keys())
	node_name = node_name_list[node_id]
	osd_idx = len(g_node_list[node_name])
	osd_name = get_osd_name(osd_idx)
	g_node_list[node_name].append(osd_name)
	return osd_name

def add_given_disk(node_id, osd_name):
	global g_node_list
	node_name_list = list(g_node_list.keys())
	node_name = node_name_list[node_id]
	g_node_list[node_name].append(osd_name)

def remove_disk_by_name(node_name, osd_name):
	global g_node_list
	g_node_list[node_name].remove(osd_name)
	if (node_name, osd_name) in g_osd_weight:
		g_osd_weight.pop((node_name, osd_name))
	return osd_name

def remove_disk(node_id, osd_idx):
	global g_node_list
	node_name_list = list(g_node_list.keys())
	node_name = node_name_list[node_id]
	osd_name = g_node_list[node_name].pop(osd_idx)
	remove_disk_by_name(node_name, osd_name)
	return osd_name

def add_node(nr_disk):
	global g_node_list
	osd_list = []
	for osd_id in range(nr_disk):
		osd_name = get_osd_name(osd_id)
		osd_list.append(osd_name)
	nr = len(g_node_list)
	node_name = get_node_name(nr)
	g_node_list[node_name] = osd_list
	return node_name

def remove_node(node_name):
	global g_node_list
	if node_name in g_node_list:
		g_node_list.pop(node_name)
		if node_name in g_node_weight:
			g_node_weight.pop(node_name)
	else:
		print('the node %s not found' % node_name)
	return

def calc_hash(s):
	m = hashlib.md5()
	m.update(s.encode('utf-8'))
	ret, = struct.unpack('Q', m.digest()[:8])
	return ret

def calc_hval(x, weight):
	a = (1<<50) * math.log2(1.0 + x) / weight
	return int(a)

def choose_node_list(pg_id):
	node_list = []
	for node_name in g_node_list:
		s = '%s-%d-%s' % (g_cluster_id, pg_id, node_name)
		x = calc_hash(s)
		if node_name in g_node_weight:
			weight = g_node_weight[node_name]
		else:
			weight = 1.0
		#hval = abs(hash((g_cluster_id, pg_id, node_name))) 
		#hval = 
		hval = calc_hval(x, weight)
		node_list.append((hval, node_name))
	node_list.sort(reverse=True)

	ret = []
	for i in range(nr_replica):
		hval, node_name = node_list[i]
		ret.append(node_name)

	return ret

def choose_osd(pg_id, node_name): 
	osd_list = g_node_list[node_name]
	idx = 0
	chosen_list = []
	for osd_name in osd_list:
		s = '%s-%d-%s-%s' % (g_cluster_id, pg_id, node_name, osd_name) 
		x = calc_hash(s)
		#hval = abs(hash((g_cluster_id, pg_id, node_name, osd_name)))
		if (node_name, osd_name) in g_osd_weight:
			weight = g_osd_weight[(node_name, osd_name)]
		else:
			weight = 1.0
		hval = calc_hval(x, weight)
		chosen_list.append((hval, osd_name))
		idx += 1
	chosen_list.sort(reverse=True)
	return chosen_list[0][1]

def crush(): 
	pg_map = {}
	for pg_id in range(nr_pg):
		osd_list = []
		node_list = choose_node_list(pg_id)
		for node_name in node_list:
			osd_name = choose_osd(pg_id, node_name)
			osd_list.append((node_name, osd_name))
		pg_map[pg_id] = osd_list
	return pg_map

def calc_stat_node(pg_map):
	stat = {}
	for i in range(nr_pg):
		for node, osd in pg_map[i]: 
			if node not in stat:
				stat[node] = 1
			else:
				stat[node] += 1
	stat_list = []
	for node in stat:
		stat_list.append((stat[node], node))
	stat_list.sort()
	return stat_list

def calc_stat_osd(pg_map):
	stat = {}
	for i in range(nr_pg):
		for node_osd in pg_map[i]: 
			if node_osd not in stat:
				stat[node_osd] = 1
			else:
				stat[node_osd] += 1
	stat_list = []
	for node_osd in stat:
		stat_list.append((stat[node_osd], node_osd))
	stat_list.sort()
	return stat_list

def calc_gini(pgnr_list):
	nr = len(pgnr_list)
	total = sum(pgnr_list)
	A_B = 0.
	B = 0.
	avg = total / nr
	cdf = 0.0
	for i in range(nr):
		cdf += avg
		A_B += cdf
	cdf = 0.0
	for i in range(nr):
		cdf += pgnr_list[i]
		B += cdf
	A = A_B - B
	return A / A_B

def calc_variance(pgnr_list):
	nr = len(pgnr_list)
	total = sum(pgnr_list)
	avg = 1.0*total / nr
	ret = 0
	for i in pgnr_list:
		a = 1.0*abs(i-avg)
		ret += a*a
	return ret/(nr-1), avg

def show_gini(pg_map):
	stat_list = calc_stat_osd(pg_map)
	pgnr_list = [i for i, j in stat_list]
	gini = calc_gini(pgnr_list)
	var, avg = calc_variance(pgnr_list)
	fmt = 'nr_pg=%5d, nr_osd=%2d, nr_node=%3d, gini = %4.2f, '\
		'var = %9.2f, std=%8.2f, avg=%6.1f, max=%5d, min=%5d'
	print(fmt % 
		(nr_pg, nr_ssd_per_node, nr_node, gini, var, math.sqrt(var), 
		avg, pgnr_list[-1], pgnr_list[0]))

def run_once():
	sys_init()
	pg_map = crush()
	show_gini(pg_map)

def test_equilibrium():
	global nr_node, nr_ssd_per_node, nr_pg

	PG_lst = [64, 256, 1024, 4096, 16384]
	NODE_lst = [3, 6, 12, 36, 64]
	OSD_lst = [2, 4, 6, 12, 24] 

	PG_dft = 1024
	NODE_dft = 6
	OSD_dft = 6

	nr_node, nr_ssd_per_node, nr_pg = NODE_dft, OSD_dft, PG_dft
	for i in PG_lst: 
		nr_pg = i
		run_once()

	nr_node, nr_ssd_per_node, nr_pg = NODE_dft, OSD_dft, PG_dft
	for i in NODE_lst: 
		nr_node = i
		run_once()

	nr_node, nr_ssd_per_node, nr_pg = NODE_dft, OSD_dft, PG_dft
	for i in OSD_lst: 
		nr_ssd_per_node = i
		run_once()

def pgmap_osdpg(pg_map):
	osd_pg = {}
	for i in pg_map:
		osd_list = pg_map[i]
		for key in osd_list:
			if key not in osd_pg:
				osd_pg[key] = []
			osd_pg[key].append(i)
	for key in osd_pg:
		osd_pg[key].sort()
	return osd_pg

def show_migrated_pg(pg_map1, pg_map2):
	nr = len(pg_map1)
	n = 0
	for i in range(nr):
		a = set(pg_map1[i])
		b = set(pg_map2[i])
		ab = list(a ^ b)
		if ab == []:
			continue
		_from = []
		for rep in pg_map1[i]:
			if rep in ab:
				_from.append(rep)
		_to = []
		for rep in pg_map2[i]:
			if rep in ab:
				_to.append(rep)

		if len(_from) != len(_to):
			print('Error')
			sys.exit(1)

		nn = len(_from)
		for j in range(nn):
			print("%s: %s ==> %s" % (i, _from[j], _to[j]))
			n += 1
	print("num_of_mig = %d" % n)

def get_line():
	line = input('cs> ')
	return line

def cmd_help(args):
	print('init     -- initialize the cluster')
	print('args     -- show or change num_pg, num_node, num_osd')
	print('crush    -- recalculate pgmap')
	print('pgmap    -- show PG map')
	print('node_stat -- show how many PG-replca on each node')
	print('osd_stat -- show how many PG-replca on each OSD')
	print('mig      -- show those migratied pg-replicas between two PG maps')
	print('gini     -- show gini-rate of a PG map')
	print('node_list -- show all node')
	print('node_add -- add a node to cluster')
	print('node_del -- remove a node from cluster')
	print('osd_list -- show all OSDs on a node')
	print('osd_add  -- add a OSD to a node')
	print('osd_del  -- remove a OSD from a node')
	print('node_wei -- show or change node weight')
	print('osd_wei  -- show or change OSD weight')

def cmd_args(args):
	global nr_node, nr_ssd_per_node, nr_pg
	print('num_pg = %d, num_node = %d, num_osd = %d' % 
		(nr_pg, nr_node, nr_ssd_per_node))
	print('Usage: args pg <pg_num>')
	print('Usage: args node <node_num>')
	print('Usage: args osd <osd_num>')
	if len(args)==3:
		if args[1]=='pg':
			nr_pg = int(args[2])
		elif args[1]=='node':
			nr_node = int(args[2])
		elif args[1]=='osd':
			nr_ssd_per_node = int(args[2])
		else:
			print('unknow args %s' % args[1])
	print('num_pg = %d, num_node = %d, num_osd = %d' % 
		(nr_pg, nr_node, nr_ssd_per_node))
	return

def cmd_gini(args, pgmap, epoch):
	if len(args)!=2:
		print('Avalible epoches are 1 - %d' % (epoch-1))
		print('Usage: gini <epoch>');
		return

	ep = int(args[1])
	if ep not in pgmap:
		print('Avalible epoches are 1 - %d' % (epoch-1))
	else:
		show_gini(pgmap[ep])
	return

def cmd_pgmap(args, pgmap, epoch):
	if len(args)!=2:
		print('Avalible epoches are 1 - %d' % (epoch-1))
		print('Usage: pgmap <epoch>');
		return

	ep = int(args[1])
	if ep not in pgmap:
		print('Avalible epoches are 1 - %d' % (epoch-1))
		return

	for i in pgmap[ep]:
		print(i, pgmap[ep][i])
	return

def cmd_stat_osd(args, pgmap, epoch):
	if len(args)!=2:
		print('Usage: stat_osd <epoch>');
		return

	ep = int(args[1])
	if ep not in pgmap:
		print('Avalible epoches are 1 - %d' % (epoch-1))
		return
	stat = calc_stat_osd(pgmap[ep])
	for i in stat:
		print(i)
	return

def cmd_stat_node(args, pgmap, epoch):
	if len(args)!=2:
		print('Usage: stat_node <epoch>');
		return

	ep = int(args[1])
	if ep not in pgmap:
		print('Avalible epoches are 1 - %d' % (epoch-1))
		return
	stat = calc_stat_node(pgmap[ep])
	for i in stat:
		print(i)
	return

def cmd_mig(args, pgmap, epoch):
	if len(args)!=3:
		print('Usage: mig <epoch1> <epoch2>')
		return

	ep1 = int(args[1])
	ep2 = int(args[2])
	if (ep1 not in pgmap) or (ep2 not in pgmap):
		print('Avalible epoches are 1 - %d' % (epoch-1))
		return

	show_migrated_pg(pgmap[ep1], pgmap[ep2])
	return

def cmd_node_list(args):
	global g_node_list
	nodes = list(g_node_list.keys())
	i = 0
	for n in nodes:
		print('node_id: %d, node_name: %s, num_osd: %d' 
			% (i, n, len(g_node_list[n])))
		i += 1
	return

def cmd_node_add(args, pgmap, epoch):
	if len(args)==2:
		nr_disk = int(args[1])
		if nr_disk<1 or nr_disk>64:
			nr_disk = OSD_dft
	else:
		nr_disk = OSD_dft
	node_name = add_node(nr_disk)
	print('osd name = %s, nr_disk = %d' % (node_name, nr_disk))
	return

def cmd_node_del(args, pgmap, epoch):
	if len(args) != 2:
		print('Usage: node_del <node_name>')
	else:
		remove_node(args[1])
	return

def cmd_osd_list(args):
	if len(args) != 2:
		print('Usage: osd_list <node_name>')
		return
	if args[1] not in g_node_list:
		print('node "%s" not found' % args[1])
		return

	osd_list = g_node_list[args[1]]
	i = 0
	for osd in osd_list:
		print('osd_id: %d, osd_name: %s' % (i, osd))
		i += 1
	return

def cmd_osd_add(args, pgmap, epoch):
	if len(args) != 2:
		print('Usagge: osd_add <node_name>')
		return

	if args[1] not in g_node_list:
		print('node "%s" not found' % args[1])
		return

	osd_name = add_disk_by_name(args[1])
	print("osd '%s' has been added to node '%s'" %
		(osd_name, args[1]))
	return

def cmd_osd_del(args, pgmap, epoch):
	if len(args) != 3:
		print('Usage: osd_del <node_name> <osd_name>')
	elif args[1] not in g_node_list:
		print('node "%s" not found' % args[1])
	elif args[2] not in g_node_list[args[1]]:
		print('osd "%s" not found' % args[1])
	else:
		osd_name = remove_disk_by_name(args[1], args[2])
		print("osd '%s' on node '%s' has been removed'" %
			(osd_name, args[1]))
	return

def cmd_node_wei_auto(args, pgmap, epoch):
	ep = int(args[2])
	if ep <= 0 or ep >= epoch:
		print('ep(%d) must be in (0, %d)' % (ep, epoch))
		return
	stat = calc_stat_node(pgmap[ep])
	nr_node = len(g_node_list)
	nr_pg = len(pgmap[ep])
	nr_pg_replica = nr_replica * nr_pg
	nr_pg_node = nr_pg_replica / nr_node
	for pg_num, node_name in stat:
		#print(pg_num, node_name, pg_num/nr_pg_node)
		g_node_weight[node_name] = nr_pg_node / pg_num
	
	
def cmd_node_wei(args, pgmap, epoch):
	if len(args)==1:
		for node_name in g_node_weight:
			print('node_name: %s, weigh: %4.2f' %
				(node_name, g_node_weight[node_name]))
	elif len(args)==3:
		if args[1]=='auto':
			cmd_node_wei_auto(args, pgmap, epoch)
			return
		weight = float(args[2])
		if args[1] not in g_node_list:
			print('node "%s" not found' % args[1])
		elif weight < 0.01 or weight > 10000.:
			print('weight(%4.2f) is not valid' %
				weight)
		else:
			g_node_weight[args[1]] = weight
	else:
		print('Usage: node_wei')
		print('       node_wei <node_name> <weigh>')
	return

def cmd_osd_wei_auto(args, pgmap, epoch):
	ep = int(args[2])
	if ep <= 0 or ep >= epoch:
		print('ep(%d) must be in (0, %d)' % (ep, epoch))
		return
	stat = calc_stat_node(pgmap[ep])
	node_stat = {}
	for pg_nr, node_name in stat:
		node_stat[node_name] = pg_nr

	osd_stat = calc_stat_osd(pgmap[ep])
	for pg_num, (node_name, osd_name) in osd_stat:
		#print(pg_num, node_name, pg_num/nr_pg_node)
		node_pg_nr = node_stat[node_name]
		node_osd_nr = len(g_node_list[node_name])
		if node_osd_nr == 0:
			continue
		avg = node_pg_nr / node_osd_nr
		g_osd_weight[(node_name, osd_name)] = avg / pg_num
	
	
def cmd_osd_wei(args, pgmap, epoch):
	if len(args)==1:
		for node_name, osd_name in g_osd_weight:
			print('(%s, %s): %4.2f' % 
				(node_name, osd_name, g_osd_weight[(node_name, osd_name)]))
	elif len(args)==3 and args[1]=='auto':
			cmd_osd_wei_auto(args, pgmap, epoch)
			return
	elif len(args)==4:
		weight = float(args[3])
		if args[1] not in g_node_list:
			print('node "%s" not found' % args[1])
		elif args[2] not in g_node_list[args[1]]:
			print('OSD "%s" not found' % args[2])
		elif weight < 0.01 or weight > 10000.:
			print('weight(%4.2f) is not valid' %
				weight)
		else:
			g_osd_weight[(args[1], args[2])] = weight
	else:
		print('Usage: osd_wei')
		print('       osd_wei <node_name> <osd_name> <weigh>')
	return

def cli():

	pgmap = {}
	epoch = 1

	while True:
		line = get_line()
		args = line.split()
		if args == []:
			continue
		elif args[0]=='quit' or args[0]=='q' or args[0]=='Q':
			break
		elif args[0]=='help' or args[0]=='h':
			cmd_help(args)
			continue
		elif args[0]=='init':
			sys_init()
			pgmap = {}
			epoch = 1
			print('cluster id = %s' % g_cluster_id)
			continue
		elif args[0]=='args':
			cmd_args(args)
			continue

		if g_cluster_id=='':
			print('must initialize the cluster')
			continue

		if args[0]=='crush':
			pgmap[epoch] = crush()
			print('pgmap epoch = %d' % epoch)
			show_gini(pgmap[epoch])
			epoch += 1
		elif args[0]=='gini':
			cmd_gini(args, pgmap, epoch)
		elif args[0]=='pgmap':
			cmd_pgmap(args, pgmap, epoch)
		elif args[0]=='osd_stat':
			cmd_stat_osd(args, pgmap, epoch)
		elif args[0]=='node_stat':
			cmd_stat_node(args, pgmap, epoch)
		elif args[0]=='mig':
			cmd_mig(args, pgmap, epoch)
		elif args[0]=='node_list': 
			cmd_node_list(args)
		elif args[0]=='node_add': 
			cmd_node_add(args, pgmap, epoch)
		elif args[0]=='node_del': 
			cmd_node_del(args, pgmap, epoch)
		elif args[0]=='osd_list':
			cmd_osd_list(args)
		elif args[0]=='osd_add': 
			cmd_osd_add(args, pgmap, epoch)
		elif args[0]=='osd_del':
			cmd_osd_del(args, pgmap, epoch)
		elif args[0]=='node_wei':
			cmd_node_wei(args, pgmap, epoch)
		elif args[0]=='osd_wei':
			cmd_osd_wei(args, pgmap, epoch)
		else:
			print("unknow command '%s'" % args[0])

def main():
	cli()

if __name__ == "__main__":
	main()


